package org.sample.algorithm.tree.rbtree;

import static org.sample.algorithm.tree.rbtree.Node.BLACK;
import static org.sample.algorithm.tree.rbtree.Node.RED;

/**
 * Created by alan on 2017/7/19.
 */
public class RBTree<T> {

    private int size = 0;

    private int modCount = 0;

    private Node<T> root;

    public Node<T> get(int key) {
        Node<T> p = root;
        while (p != null) {
            if (key < p.key) {
                p = p.left;
            } else if (key > p.key) {
                p = p.right;
            } else {
                return p;
            }
        }
        return null;
    }

    public T put(int key, T value) {
        Node<T> t = root;
        if (t == null) {
            root = new Node<>(key, value, null);
            size = 1;
            modCount++;
            return null;
        }
        Node<T> parent;
        do {
            parent = t;
            if (key < t.key) {
                t = t.left;
            } else if (key > t.key) {
                t = t.right;
            } else {
                return t.setValue(value);
            }
        } while (t != null);

        Node<T> node = new Node<>(key, value, parent);
        if (key < parent.key) {
            parent.left = node;
        } else {
            parent.right = node;
        }
        fixAfterInsertion(node);
        size++;
        modCount++;
        return null;
    }

    private void fixAfterInsertion(Node<T> x) {
        x.color = RED;

        while (x != null && x != root && x.parent.color == RED) {
            if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
                Node<T> y = rightOf(parentOf(parentOf(x)));
                if (colorOf(y) == RED) {
                    setColor(parentOf(x), BLACK);
                    setColor(y, BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    x = parentOf(parentOf(x));
                } else {
                    if (x == rightOf(parentOf(x))) {
                        x = parentOf(x);
                        rotateLeft(x);
                    }
                    setColor(parentOf(x), BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    rotateRight(parentOf(parentOf(x)));
                }
            } else {
                Node<T> y = leftOf(parentOf(parentOf(x)));
                if (colorOf(y) == RED) {
                    setColor(parentOf(x), BLACK);
                    setColor(y, BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    x = parentOf(parentOf(x));
                } else {
                    if (x == leftOf(parentOf(x))) {
                        x = parentOf(x);
                        rotateRight(x);
                    }
                    setColor(parentOf(x), BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    rotateLeft(parentOf(parentOf(x)));
                }
            }
        }
        root.color = BLACK;
    }


    private void rotateLeft(Node<T> p) {
        if (p != null) {
            Node<T> r = p.right;
            p.right = r.left;
            if (r.left != null)
                r.left.parent = p;
            r.parent = p.parent;
            if (p.parent == null)
                root = r;
            else if (p.parent.left == p)
                p.parent.left = r;
            else
                p.parent.right = r;
            r.left = p;
            p.parent = r;
        }
    }

    private void rotateRight(Node<T> p) {
        if (p != null) {
            Node<T> l = p.left;
            p.left = l.right;
            if (l.right != null)
                l.right.parent = p;
            l.parent = p.parent;
            if (p.parent == null)
                root = l;
            else if (p.parent.right == p)
                p.parent.right = l;
            else
                p.parent.left = l;
            l.right = p;
            p.parent = l;
        }
    }

    private static <T> Node<T> parentOf(Node<T> x) {
        return x == null ? null : x.parent;
    }

    private static <T> Node<T> leftOf(Node<T> x) {
        return x == null ? null : x.left;
    }

    private static <T> Node<T> rightOf(Node<T> x) {
        return x == null ? null : x.right;
    }

    private static <T> boolean colorOf(Node<T> x) {
        return x == null ? BLACK : x.color;
    }

    private static <T> void setColor(Node<T> x, boolean color) {
        if (x != null)
            x.color = color;
    }


    public void show() {
        show(root);
    }

    private static void show(Node<?> root) {
        if (root != null) {
            show(root.left);
            System.out.println(root);
            show(root.right);
        }
    }
}